package com.xialuo.study.nodes;

import com.xialuo.study.state.ReactState;
import org.bsc.langgraph4j.action.AsyncNodeAction;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.model.function.FunctionCallback;
import org.springframework.ai.tool.resolution.ToolCallbackResolver;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

/**
 * 处理工具执行的节点
 */
public class ToolNode implements AsyncNodeAction<ReactState> {
    private final List<FunctionCallback> tools;
    private final ToolCallbackResolver resolver;

    public ToolNode(List<FunctionCallback> tools) {
        this.tools = tools;
        this.resolver = null;
    }

    public ToolNode(ToolCallbackResolver resolver) {
        this.resolver = resolver;
        this.tools = null;
    }

    @Override
    public CompletableFuture<Map<String, Object>> apply(ReactState state) {
        List<Message> messages = (List<Message>) state
                .value("messages")
                .get();
        AssistantMessage lastAssistantMessage = (AssistantMessage) messages.get(messages.size() - 1);

        // 处理工具调用
        if (lastAssistantMessage.hasToolCalls()) {
            List<Message> toolResponseMessages = new ArrayList<>();

            lastAssistantMessage
                    .getToolCalls()
                    .forEach(toolCall -> {
                        String toolName = toolCall.name();
                        String toolInput = toolCall.arguments();
                        String toolOutput;

                        if (resolver != null) {
                            // 使用解析器查找并执行工具
                            toolOutput = resolver
                                    .resolve(toolName)
                                    .call(toolInput);
                        } else {
                            // 在列表中查找工具并执行
                            FunctionCallback tool = tools
                                    .stream()
                                    .filter(t -> t
                                            .getName()
                                            .equals(toolName))
                                    .findFirst()
                                    .orElseThrow(() -> new IllegalArgumentException("找不到工具: " + toolName));

                            toolOutput = tool.call(toolInput);
                        }

                        // 创建工具响应消息
                        ToolResponseMessage toolResponseMessage = new ToolResponseMessage(List.of(new ToolResponseMessage.ToolResponse(toolCall.id(), toolName, toolOutput)));
                        toolResponseMessages.add(toolResponseMessage);
                    });

            // 将工具响应添加到消息列表
            messages.addAll(toolResponseMessages);
        }

        // 返回更新后的状态
        Map<String, Object> newState = new HashMap<>();
        newState.put("messages", messages);
        return CompletableFuture.completedFuture(newState);
    }
}
